package prod_cons;

import java.util.Random;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;

public class WorkStealingContainer implements TaskContainer {
	int consumersNum;
	ConcurrentLinkedQueue<Runnable>[] queues;
	
	@SuppressWarnings("unchecked")
	public void setConsumersNum(int consumersNum) {
		this.consumersNum = consumersNum;
		queues = new ConcurrentLinkedQueue[consumersNum];
	
		for(int i = 0; i < consumersNum; i++) {
			queues[i] = new ConcurrentLinkedQueue<Runnable>();
		}
	}
	
	private static final ThreadLocal<Integer> queueNum = new ThreadLocal<Integer>() {
		protected Integer initialValue() {
			return new Integer(0);
		}
	};
	
	AtomicInteger consumerCount = new AtomicInteger(0);
	public void registerConsumer() {
		queueNum.set(consumerCount.getAndIncrement());
	}
	
	Random random = new Random();
	
	@Override
	public void put(Runnable task) {
		// put to a random queue
		queues[random.nextInt(this.consumersNum)].offer(task);
	}

	@Override
	public Runnable get() {
		int myQueue = queueNum.get();
		Runnable task = queues[myQueue].poll();
		if (task != null) return task;
		for(int i = 1; i < this.consumersNum; i++) {
			task = queues[(myQueue+i)%consumersNum].poll();
			if (task != null) return task;
		}
		return null;
	}
}
